package com.custom.source;

import lombok.val;
import lombok.var;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;

import javax.annotation.Nullable;
import java.lang.reflect.ParameterizedType;
import java.util.ArrayDeque;
import java.util.Collection;
import java.util.List;
import java.util.Queue;

/**
 * 分片枚举器
 *
 * @param <SplitT>   分片类型
 * @param <EnumChkT> 枚举分片保存点类型
 */
public abstract class ClinkSplitEnumerator<SplitT extends ClinkSourceSplit, EnumChkT extends ClinkEnumCheckpoint<SplitT>> implements SplitEnumerator<SplitT, EnumChkT> {

    protected final Queue<SplitT> pendingSplits = new ArrayDeque<>();
    protected SplitEnumeratorContext<SplitT> context;

    public ClinkSplitEnumerator(SplitEnumeratorContext<SplitT> context) {
        this.context = context;
    }

    public ClinkSplitEnumerator(SplitEnumeratorContext<SplitT> context, EnumChkT checkpoint) {
        this(context);
        pendingSplits.addAll(checkpoint.getState());
    }

    /**
     * 如果还有未分配的队列,则分配出去
     */
    @Override
    public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
        val split = pendingSplits.poll();
        if (split != null) {
            context.assignSplit(split, subtaskId);
        } else {
            context.signalNoMoreSplits(subtaskId);
        }
    }

    /**
     * 将分配失败split返回到未处理分片队列中
     */
    @Override
    public void addSplitsBack(List<SplitT> splits, int subtaskId) {
        pendingSplits.addAll(splits);
    }

    /**
     * 将未分配的分片保存快照
     */
    @Override
    public EnumChkT snapshotState(long checkpointId) throws Exception {
        var type = this.getClass().getGenericSuperclass();
        if (type instanceof ParameterizedType) {
            val parameterizedType = (ParameterizedType) type;
            // 获取当前抽象类第二个泛型class对象
            type = parameterizedType.getActualTypeArguments()[1];
            if (type instanceof Class<?>) {
                val clazz = (Class<?>) type;
                val constructor = clazz.getConstructor(Collection.class);
                @SuppressWarnings("unchecked")
                val checkpoint = (EnumChkT) constructor.newInstance(pendingSplits);
                return checkpoint;
            }
        }

        throw new UnsupportedOperationException("can not snapshot checkpoint");
    }
}
